import torch
import torchvision
import torch.nn as nn

class ResNet(torch.nn.Module):
    def __init__(self, classNum):
        super(ResNet, self).__init__()
        self.resnet = torchvision.models.resnet18()
        self.resnet.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3,
                                 bias=False)
        self.resnet.fc = nn.Linear(in_features=512, out_features=classNum)
        print('cur conv is resnet18')

    def forward(self, x):
        # x shape = (batch, channel, w, h)
        x = self.resnet(x)
        return x

class VGGNet(torch.nn.Module):
    def __init__(self, classNum):
        super(VGGNet, self).__init__()
        self.vgg16 = torchvision.models.vgg16_bn(num_classes=classNum)
        self.vgg16.features._modules['0'] = nn.Conv2d(1, 64, kernel_size=3, padding=1)
        self.vgg16.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.vgg16.classifier = nn.Sequential(
            nn.Linear(512, 1024),
            nn.Linear(1024, classNum),
        )
        # self.vgg16.classifier[0] = nn.Linear(512, 4096)

        print('cur conv is vgg16')


    def forward(self, x):
        x = self.vgg16(x)
        return x

class AlexNet(nn.Module):

    def __init__(self, num_classes=1000):
        super(AlexNet, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 64, kernel_size=11, stride=4, padding=2),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(64, 192, kernel_size=5, padding=2),
            nn.BatchNorm2d(192),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(192, 384, kernel_size=3, padding=1),
            nn.BatchNorm2d(384),
            nn.ReLU(inplace=True),
            nn.Conv2d(384, 256, kernel_size=3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
        )
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.classifier = nn.Sequential(
            nn.Dropout(),
            nn.Linear(256 * 6 * 6, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(inplace=True),
            nn.Linear(4096, num_classes),
        )
        self.fc = nn.Linear(256, num_classes)

    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.fc(x)
        return x

class AlexNetUltra(torch.nn.Module):
    def __init__(self, classNum):
        super(AlexNetUltra, self).__init__()
        self.alexnet = AlexNet(num_classes=classNum)
        # self.alexnet.features[0] = nn.Conv2d(in_channels=1, out_channels=64, kernel_size=11, stride=4, padding=2)
        print('cur conv is alexnet')

    def forward(self, x):
        x = self.alexnet(x)
        return x

class MobileNetv2(torch.nn.Module):
    def __init__(self, classNum):
        super(MobileNetv2, self).__init__()
        self.mobilenet = torchvision.models.mobilenet_v2(num_classes=classNum)
        self.mobilenet.features[0][0] = torch.nn.Conv2d(1, 32, kernel_size=3, stride=2, padding=1, bias=False)
        print('cur conv is mobilenet')
        # print(self.mobilenet)

    def forward(self, x):
        x = self.mobilenet(x)
        return x